import torch
from torch.autograd import Variable as V

import cv2
import os
import importlib
import numpy as np
from tqdm import tqdm

SOURCE_PATH = 'H:/Teach/data2/test'
TARGET_PATH = 'H:/Teach/data2/result_unet'
MODEL_FILENAME = 'H:/Teach/data2/model/UNet_page0/model000000200.model'

MODUAL_NAME = 'models.bilstm'
MODEL_NAME = 'BiLSTM'


def test_one_img_from_path_4(net, filename):
    readFileA = np.fromfile(filename, dtype=np.float32)
    sagmentA = readFileA[0: len(readFileA)]
    imgA = np.reshape(sagmentA, (256, 1024))
    imgA = torch.FloatTensor(imgA)

    # 取对数
    imgA = torch.log10(imgA)
    # Amax = torch.max(imgA)
    # Amin = torch.min(imgA)
    Amin = -2.0
    Amax = 4.0
    # 映射
    imgA = (imgA - Amin) / (Amax - Amin)
    imgA = imgA.unsqueeze(0).unsqueeze(0)

    with torch.no_grad():
        mask = net.forward(imgA).squeeze().cpu().data.numpy()

    return mask

if __name__ == "__main__":
    image_list = list(filter(lambda x: x.find('.bmdl2') != -1, os.listdir(SOURCE_PATH)))

    if not os.path.exists(TARGET_PATH):
        os.mkdir(TARGET_PATH)

    net = getattr(importlib.import_module(MODUAL_NAME), MODEL_NAME)
    net = net().cuda()
    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
    net.load_state_dict(torch.load(MODEL_FILENAME))
    net.eval()

    pbar = tqdm(total=len(image_list))  # 进度条

    for i, name in enumerate(image_list):
        mask = test_one_img_from_path_4(net, os.path.join(SOURCE_PATH, name))
        mask = cv2.blur(mask, (11, 11))
        mask = cv2.resize(mask, (128, 32), interpolation=cv2.INTER_CUBIC)
        fname = os.path.join(TARGET_PATH, name[:-5])
        #np.save(fname, mask)
        np.savetxt(fname+'.txt', mask, fmt="%.8f")
        pbar.update(1)

    pbar.close()